// https://leetcode.cn/problems/top-k-frequent-elements/
package leet347topkfrequentelement

import "container/heap"

// 思路：
// 1. 使用一个map统计每个元素出现的频率
// 2. 使用一个最小堆存储频率最高的k个元素
// 3. 遍历map，将频率和元素存入最小堆中
// 4. 如果堆的大小超过k，则弹出堆顶元素
// 5. 最后返回堆中的元素
// 6. 时间复杂度O(nlogk)，空间复杂度O(n)

// 为什么使用堆排序会比快速排序快？
// 1. 堆排序的时间复杂度是O(nlogn)，而快速排序的平均时间复杂度是O(nlogn)，最坏时间复杂度是O(n^2)
// 2. 堆排序的空间复杂度是O(1)，而快速排序的空间复杂度是O(logn)
// 3. 堆排序是一个不稳定的排序算法，而快速排序是一个不稳定的排序算法
// 4. 堆排序的实现比较简单，而快速排序的实现比较复杂
// 5. 堆排序的性能受数据分布的影响较小，而快速排序的性能受数据分布的影响较大

// 定义一个堆元素结构体
type HeapElement struct {
	Val   int // 元素值
	Freq  int // 元素频率
	Index int // 元素在堆中的索引
}

// 定义一个优先队列实现 heap.Interface 接口，并持有堆元素
type MinHeap []*HeapElement

func (h MinHeap) Len() int {
	return len(h)
}
func (h MinHeap) Less(i, j int) bool {
	return h[i].Freq < h[j].Freq
}
func (h MinHeap) Swap(i, j int) {
	h[i], h[j] = h[j], h[i]
	// 更新索引的目的是为了在堆中保持元素的正确位置
	// 因为在交换元素时，索引也需要交换
	// 能不能不用Index这个属性呢？@todo
	h[i].Index = i
	h[j].Index = j
}
func (h *MinHeap) Push(x interface{}) {
	item := x.(*HeapElement)
	item.Index = len(*h)
	*h = append(*h, item)
}

// Pop 方法从堆中弹出最小元素
// 注意：Pop 方法会修改堆的结构，所以需要实现 heap.Interface 接口
// 堆会自动调整结构
// 1. 将最后一个元素和堆顶的元素交换
// 2. 然后将最后一个元素的索引设置为-1，表示该元素不在堆中
// 3. 然后将堆的大小减1
// 4. 最后返回弹出的元素
func (h *MinHeap) Pop() interface{} {
	old := *h
	n := len(old)
	x := old[n-1]
	old[n-1] = nil // avoid memory leak
	x.Index = -1   // for safety
	*h = old[0 : n-1]
	return x
}
func (h *MinHeap) Top() *HeapElement {
	if h.Len() == 0 {
		return nil
	}
	return (*h)[0]
}

// TopKFrequentElements 函数返回出现频率最高的k个元素
func TopKFrequentElements(nums []int, k int) []int {
	// 1. 使用一个map统计每个元素出现的频率
	freqMap := make(map[int]int)
	for _, num := range nums {
		freqMap[num]++
	}

	// 2. 使用一个最小堆存储频率最高的k个元素
	minHeap := &MinHeap{}
	for num, freq := range freqMap {
		heap.Push(minHeap, &HeapElement{Val: num, Freq: freq})
		// 如果堆的大小超过k，则弹出堆顶元素
		// 堆顶元素是频率最低的元素
		if minHeap.Len() > k {
			heap.Pop(minHeap)
		}
	}

	// 3. 返回堆中的元素
	result := make([]int, 0, k)
	for minHeap.Len() > 0 {
		result = append(result, heap.Pop(minHeap).(*HeapElement).Val)
	}
	return result
}
